

import os
import cv2
import math
import numpy as np
import mindspore as ms
import mindspore.nn as nn
import mindspore.dataset.vision as dv


# 设置成静态图
ms.set_context(mode=ms.PYNATIVE_MODE)

class SPPLayer(nn.Cell):
    def __init__(self, num_layers, pool_type='max_pool'):
        super(SPPLayer, self).__init__()

        self.num_levels = num_layers
        self.pool_type = pool_type
        self.flatten = nn.Flatten()
        self.cat = ms.ops.Concat(axis=1)

    def construct(self, input_feature):
        num, c, h, w = input_feature.shape
        for i in range(self.num_levels):
            level = i + 1
            kernel_size = (math.ceil(h / level), math.ceil(w / level))
            stride = (math.ceil(h / level), math.ceil(w / level))
            pooling = (math.floor((kernel_size[1] * level - w + 1) / 2),
                       math.floor((kernel_size[1] * level - w + 1) / 2),
                       math.floor((kernel_size[0] * level - h + 1) / 2),
                       math.floor((kernel_size[0] * level - h + 1) / 2))
            # 选择池化方式
            pad = nn.ReflectionPad2d(pooling)
            if self.pool_type == 'max_pool':
                cov = nn.MaxPool2d(kernel_size=kernel_size, stride=stride)
                tensor = cov(pad(input_feature))
            else:
                cov = nn.AvgPool2d(kernel_size=kernel_size, stride=stride)
                tensor = cov(pad(input_feature))

            # 展开、拼接
            if i == 0:
                x_flatten = self.flatten(tensor)
            else:
                x_flatten = self.cat((x_flatten, self.flatten(tensor)))

        return x_flatten


class SPPLeNet5(nn.Cell):
    """
    LeNet-5网络结构
    """
    def __init__(self, num_class=10, num_channel=1, num_layers=9):
        super(SPPLeNet5, self).__init__()
        self.conv1 = nn.Conv2d(num_channel, 6, 5, pad_mode='valid')
        self.conv2 = nn.Conv2d(6, 16, 5, pad_mode='valid')
        self.fc1 = nn.Dense(880, 4096)
        self.fc2 = nn.Dense(4096, 2048)
        self.fc3 = nn.Dense(2048, num_class)
        self.relu = nn.ReLU()
        self.max_pool2d = nn.MaxPool2d(kernel_size=2, stride=2)
        self.spp = SPPLayer(num_layers, pool_type='max_pool')
        self.flatten = nn.Flatten()

    def construct(self, x):
        # 使用定义好的运算构建前向网络
        x = self.conv1(x)
        x = self.relu(x)
        x = self.max_pool2d(x)
        x = self.conv2(x)
        x = self.relu(x)
        x = self.max_pool2d(x)
        x = self.spp(x)

        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        x = self.relu(x)
        x = self.fc3(x)
        return x


class CaltechDataset:
    """自定义数据集类"""

    def __init__(self, file_path, train=False):

        if train:
            path_train = os.path.join(file_path, 'train')
            self.data, self.label = self.__read_train(path_train)

        else:
            path_test = os.path.join(file_path, 'test')
            self.data = self.__read_test(path_test)
            self.label = np.array([-1 for _ in range(len(self.data))])

    def __getitem__(self, index):
        """自定义随机访问函数"""
        # return self.data[index], self.label[index]
        return self.data[index], self.label[index]

    def __len__(self):
        """自定义获取样本数据量函数"""
        return len(self.data)

    def __read_train(self, file_path):

        data_list = []
        label_list = []
        for label in range(256):
            file_path_temp = os.path.join(file_path, str(label+1))
            file_name_list = os.listdir(file_path_temp)
            file_name_list.sort(key=lambda _: int(_[:-4]))
            for file_name in file_name_list:
                # print(os.path.join(file_path_temp, file_name))
                img_cv = cv2.imread(os.path.join(file_path_temp, file_name))
                data_list.append(img_cv.astype(dtype=np.float32))
                label_list.append(label)
        return np.array(data_list, dtype=object), np.array(label_list, dtype=np.int32)

    def __read_test(self, file_path):

        data_list = []

        file_name_list = os.listdir(file_path)
        file_name_list.sort(key=lambda _: int(_[:-4]))
        for file_name in file_name_list:
            img_cv = cv2.imread(os.path.join(file_path, file_name))
            data_list.append(img_cv.astype(dtype=np.float32))

        return np.array(data_list)

    def test(self, file_path):
        # return self.__read_train(file_path)
        return self.data, self.label


class ImagProperty:
    def __init__(self, image_size=256, prob=0.5):
        self.image_size = image_size
        self.mean = [0.5 * 255] * 3
        self.std = [0.5 * 255] * 3
        self.prob = prob

        self.trans_train = [
            # dv.Resize((self.image_size, self.image_size)),
            dv.Normalize(mean=self.mean, std=self.std),
            dv.RandomHorizontalFlip(prob=self.prob),
            dv.HWC2CHW()
        ]

        self.trans_test = [
            # dv.Resize((self.image_size, self.image_size)),
            dv.Normalize(mean=self.mean, std=self.std),
            dv.HWC2CHW()
        ]


def gen_dataset(dataset_class, train='train', batch_size=1):
    img_tans = ImagProperty()
    if train == 'train':
        dataset = ms.dataset.GeneratorDataset(dataset_class, ['image', 'label'], shuffle=True)
        dataset = dataset.map(operations=img_tans.trans_train)
        dataset = dataset.batch(batch_size)
        return dataset

    elif train == 'train_val':
        dataset = ms.dataset.GeneratorDataset(dataset_class, ['image', 'label'], shuffle=False)
        dataset = dataset.map(operations=img_tans.trans_test)
        dataset = dataset.batch(batch_size)
        return dataset

    else:
        dataset = ms.dataset.GeneratorDataset(dataset_class, ['image', 'label'], shuffle=False)
        dataset = dataset.map(operations=img_tans.trans_test)
        dataset = dataset.batch(batch_size)
        dataset = dataset.create_dict_iterator()
        return dataset


if __name__ == '__main__':

    input_x = ms.Tensor(np.ones([10, 3, 70, 110]), dtype=ms.float32)

    net = SPPLeNet5(num_class=256, num_channel=3, num_layers=3)
    x = net(input_x)

    print(np.shape(x))